import paddle
import paddle.nn as nn
from .untils import load_pretrained


class _NonLocalBlockND(nn.Layer):
    def __init__(self, in_channels, inter_channels=None, dimension=3, sub_sample=True, bn_layer=True):
        super(_NonLocalBlockND, self).__init__()

        assert dimension in [1, 2, 3]

        self.dimension = dimension
        self.sub_sample = sub_sample

        self.in_channels = in_channels
        self.inter_channels = inter_channels

        if self.inter_channels is None:
            self.inter_channels = in_channels // 2
            if self.inter_channels == 0:
                self.inter_channels = 1

        if dimension == 3:
            conv_nd = nn.Conv3D
            max_pool_layer = nn.MaxPool3D(kernel_size=(1, 2, 2))
            bn = nn.BatchNorm3D
        elif dimension == 2:
            conv_nd = nn.Conv2D
            max_pool_layer = nn.MaxPool2D(kernel_size=(2, 2))
            bn = nn.BatchNorm2D
        else:
            conv_nd = nn.Conv1D
            max_pool_layer = nn.MaxPool1D(kernel_size=(2))
            bn = nn.BatchNorm1D

        self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
                         kernel_size=1, stride=1, padding=0)

        if bn_layer:
            self.W = nn.Sequential(
                conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
                        kernel_size=1, stride=1, padding=0),
                bn(self.in_channels)
            )

        else:
            self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,
                             kernel_size=1, stride=1, padding=0)

        self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
                             kernel_size=1, stride=1, padding=0)

        self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,
                           kernel_size=1, stride=1, padding=0)

        if sub_sample:
            self.g = nn.Sequential(self.g, max_pool_layer)
            self.phi = nn.Sequential(self.phi, max_pool_layer)

    def forward(self, x, return_nl_map=False):
        """
        :param x: (b, c, t, h, w)
        :param return_nl_map: if True return z, nl_map, else only return z.
        :return:
        """

        batch_size = x.shape[0]

        g_x = paddle.reshape(self.g(x), (batch_size, self.inter_channels, -1))
        g_x = g_x.transpose([0, 2, 1])

        theta_x = paddle.reshape(self.theta(x), (batch_size, self.inter_channels, -1))
        theta_x = theta_x.transpose([0, 2, 1])
        phi_x = paddle.reshape(self.phi(x), (batch_size, self.inter_channels, -1))
        f = paddle.matmul(theta_x, phi_x)
        N = f.shape[-1]
        f_div_C = f / N

        y = paddle.matmul(f_div_C, g_x)
        y = y.transpose([0, 2, 1])
        y = y.reshape([batch_size, self.inter_channels, *x.shape[2:]])
        W_y = self.W(y)
        z = W_y + x

        if return_nl_map:
            return z, f_div_C
        return z


class NONLocalBlock1D(_NonLocalBlockND):
    def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):
        super(NONLocalBlock1D, self).__init__(in_channels,
                                              inter_channels=inter_channels,
                                              dimension=1, sub_sample=sub_sample,
                                              bn_layer=bn_layer)


# 修改完成
class Aggregate(nn.Layer):
    def __init__(self, len_feature):
        super(Aggregate, self).__init__()
        bn = nn.BatchNorm1D
        self.len_feature = len_feature
        self.conv_1 = nn.Sequential(
            nn.Conv1D(in_channels=len_feature, out_channels=512, kernel_size=3,
                      stride=1, dilation=1, padding=1),
            nn.ReLU(),
            bn(512)
        )
        self.conv_2 = nn.Sequential(
            nn.Conv1D(in_channels=len_feature, out_channels=512, kernel_size=3,
                      stride=1, dilation=2, padding=2),
            nn.ReLU(),
            bn(512)
        )
        self.conv_3 = nn.Sequential(
            nn.Conv1D(in_channels=len_feature, out_channels=512, kernel_size=3,
                      stride=1, dilation=4, padding=4),
            nn.ReLU(),
            bn(512)
        )
        self.conv_4 = nn.Sequential(
            nn.Conv1D(in_channels=len_feature, out_channels=512, kernel_size=1,
                      stride=1, padding=0),
            nn.ReLU(),
        )
        self.conv_5 = nn.Sequential(
            nn.Conv1D(in_channels=2048, out_channels=len_feature, kernel_size=3,
                      stride=1, padding=1),  # should we keep the bias?
            nn.ReLU(),
            bn(len_feature),
        )

        self.non_local = NONLocalBlock1D(512, sub_sample=False, bn_layer=True)

    def forward(self, x):
        # x: (B, T, F)
        out = x.transpose([0, 2, 1])
        residual = out

        out1 = self.conv_1(out)
        out2 = self.conv_2(out)

        out3 = self.conv_3(out)
        out_d = paddle.concat([out1, out2, out3], axis=1)
        out = self.conv_4(out)
        out = self.non_local(out)
        out = paddle.concat((out_d, out), axis=1)
        out = self.conv_5(out)  # fuse all the features together
        out = out + residual
        out = out.transpose([0, 2, 1])
        return out


class RTFM(nn.Layer):
    def __init__(self, n_features, batch_size, crop, **kwargs):
        """
        RTFM是一个用于弱监督视频异常检测的模型。它的主要任务是区分正常和异常的视频片段。在训练过程中，RTFM使用的是弱标签视频，即视频级别的标签
        （正常或异常）。RTFM通过学习一个时间特征幅度映射函数，来检测稀有的异常片段，并保证正常和异常片段之间的较大间隔。此外，RTFM还采用膨胀卷
        积和自注意力机制来捕捉长程和短程时间依赖关系，以更准确地学习特征幅度。
        RTFM的效果提升主要归因于以下几点：
        1. RTFM能够检测出包含在正常视频中的稀有的异常片段，这在使用传统的异常检测方法时通常很难做到。
        2. RTFM通过学习一个时间特征幅度映射函数，可以更好地表示和区分正常和异常的视频片段。
        3. RTFM还采用了膨胀卷积和自注意力机制，以捕捉更广泛的时间依赖关系，从而更准确地学习特征幅度。
        这些改进使得RTFM在四个基准数据集上优于几种最先进的方法，并显著提高了对微妙异常的辨别力和样本效率。

        Args:
            n_features (int): 特征的维度，通常是2048【即视频压缩的维度】
            batch_size (int): 训练时同批的数据
        """
        super(RTFM, self).__init__()
        self.crop = crop
        self.batch_size = batch_size
        self.num_segments = 32
        self.k_abn = self.num_segments // 10
        self.k_nor = self.num_segments // 10

        self.Aggregate = Aggregate(len_feature=n_features)
        self.fc1 = nn.Linear(n_features, 512)
        self.fc2 = nn.Linear(512, 128)
        self.fc3 = nn.Linear(128, 1)

        self.drop_out = nn.Dropout(0.7)
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

        if "pretrained" in kwargs.keys() and kwargs["pretrained"] is not None:
            load_pretrained(self, kwargs["pretrained"])

    def forward(self, inputs):

        k_abn = self.k_abn
        k_nor = self.k_nor

        out = inputs
        bs, ncrops, t, f = out.shape

        out = out.reshape([-1, t, f])

        out = self.Aggregate(out)

        out = self.drop_out(out)

        features = out
        scores = self.relu(self.fc1(features))
        scores = self.drop_out(scores)
        scores = self.relu(self.fc2(scores))
        scores = self.drop_out(scores)
        scores = self.sigmoid(self.fc3(scores))

        scores = paddle.mean(scores.reshape([bs, ncrops, -1]), axis=1)
        scores = scores.unsqueeze(axis=2)

        normal_features = features[0:self.batch_size * self.crop]
        normal_scores = scores[0:self.batch_size]
        abnormal_features = features[self.batch_size * self.crop:]
        abnormal_scores = scores[self.batch_size:]
        feat_magnitudes = paddle.norm(features, p=2, axis=2)
        feat_magnitudes = paddle.mean(feat_magnitudes.reshape([bs, ncrops, -1]), axis=1)
        nfea_magnitudes = feat_magnitudes[0:self.batch_size]  # normal feature magnitudes
        afea_magnitudes = feat_magnitudes[self.batch_size:]  # abnormal feature magnitudes
        n_size = nfea_magnitudes.shape[0]

        if nfea_magnitudes.shape[0] == 1:  # this is for inference, the batch size is 1
            afea_magnitudes = nfea_magnitudes
            abnormal_scores = normal_scores
            abnormal_features = normal_features

        select_idx = paddle.ones_like(nfea_magnitudes)
        select_idx = self.drop_out(select_idx)

        #######  process abnormal videos -> select top3 feature magnitude  #######
        afea_magnitudes_drop = afea_magnitudes * select_idx
        idx_abn = paddle.topk(afea_magnitudes_drop, k_abn, axis=1)[1]
        idx_abn_feat = idx_abn.unsqueeze(2).expand([-1, -1, abnormal_features.shape[2]])

        abnormal_features = abnormal_features.reshape([n_size, ncrops, t, f])
        abnormal_features = abnormal_features.transpose([1, 0, 2, 3])

        total_select_abn_feature = paddle.zeros([0])
        for abnormal_feature in abnormal_features:
            feat_select_abn = paddle.take_along_axis(abnormal_feature, idx_abn_feat, 1)  # top 3 features magnitude in abnormal bag
            total_select_abn_feature = paddle.concat((total_select_abn_feature, feat_select_abn))

        idx_abn_score = idx_abn.unsqueeze(2).expand([-1, -1, abnormal_scores.shape[2]])
        score_abnormal = paddle.mean(paddle.take_along_axis(abnormal_scores, idx_abn_score, 1),
                                     axis=1)  # top 3 scores in abnormal bag based on the top-3 magnitude

        ####### process normal videos -> select top3 feature magnitude #######

        select_idx_normal = paddle.ones_like(nfea_magnitudes)
        select_idx_normal = self.drop_out(select_idx_normal)
        nfea_magnitudes_drop = nfea_magnitudes * select_idx_normal
        idx_normal = paddle.topk(nfea_magnitudes_drop, k_nor, axis=1)[1]
        idx_normal_feat = idx_normal.unsqueeze(2).expand([-1, -1, normal_features.shape[2]])

        normal_features = normal_features.reshape([n_size, ncrops, t, f])
        normal_features = normal_features.transpose([1, 0, 2, 3])

        total_select_nor_feature = paddle.zeros([0])
        for nor_fea in normal_features:
            feat_select_normal = paddle.take_along_axis(nor_fea, idx_normal_feat, 1)  # top 3 features magnitude in normal bag (hard negative)
            total_select_nor_feature = paddle.concat((total_select_nor_feature, feat_select_normal))

        idx_normal_score = idx_normal.unsqueeze(2).expand([-1, -1, normal_scores.shape[2]])
        score_normal = paddle.mean(paddle.take_along_axis(normal_scores, idx_normal_score, 1), axis=1)  # top 3 scores in normal bag

        feat_select_abn = total_select_abn_feature
        feat_select_normal = total_select_nor_feature

        return {
            "score_abnormal": score_abnormal,
            "score_normal": score_normal,
            "feat_select_normal": feat_select_normal,
            "feat_select_abn": feat_select_abn,
            "feat_magnitudes": feat_magnitudes,
            "scores": scores
        }


if __name__ == "__main__":
    model = RTFM(2048, 16)
    paddle.flops(model, [32, 16, 10, 2048], print_detail=True)